"""
vector_store.py是EduRAG系统的核心模块之一，封装了与Milvus向量数据库的交互逻辑。
它负责将文档转化为向量存储到数据库中，并提供高效的混合检索功能。
通过结合BGE-M3嵌入模型和重排序机制，该模块确保系统能够快速检索到与用户查询最相关的文档。
"""
import hashlib
import os.path
import sys
import torch.cuda
from langchain_core.documents import Document
# 导包
# 导入混合向量生成模型
from milvus_model.hybrid import BGEM3EmbeddingFunction
# 导入milvus相关类，用于操作Milvus向量数据库
from pymilvus import MilvusClient, DataType, WeightedRanker, AnnSearchRequest
# 导入重排序模型
from sentence_transformers import CrossEncoder

# 导入基础模块配置和日志模块
from base.config import get_config
from base.logger import get_logger
from rag_qa.core.document_processor import process_documents

# 获取当前文件所在目录的绝对路径
local_path = os.path.abspath(os.path.dirname(__file__))
print(local_path)
# 项目主目录
rag_qa_path = os.path.abspath(os.path.dirname(local_path))
print(rag_qa_path)
# 项目根目录
project_root = os.path.abspath(os.path.dirname(rag_qa_path))
print(project_root)
# 添加路径到系统路径
sys.path.insert(0, project_root)
sys.path.insert(0, rag_qa_path)

# 实例化配置对象和日志对象
config = get_config()
logger = get_logger()


class VectorStore:
    """
    VectorStore类封装了与Milvus向量数据库的交互逻辑。以及处理文档的向量化存储和检索。
    """

    def __init__(self, collection_name=config.MILVUS_COLLECTION_NAME,
                 host=config.MILVUS_HOST,
                 port=config.MILVUS_PORT,
                 database=config.MILVUS_DATABASE_NAME):
        # 1：初始化向量存储配置
        self.collection_name = collection_name
        self.host = host
        self.port = port
        self.database = database

        # 2：初始化日志模块
        self.logger = logger

        # 3：初始化模型
        # 3.1：初始化重排序模型
        # 参数说明：
        #   rerank_model_path: 重排序模型的路径，用于加载预训练的交叉编码器模型
        #   device：指定模型运行的设备，这里设置为cpu运行，以减少内存占用，默认为cuda
        #           如果在CPU运行可能会比较慢，建议在有GPU的环境下使用，还有其他设备可以选择，如"mps"（Apple Silicon）
        #           这里选择cpu运行是因为在没有GPU的情况下，也可以运行模型，只是速度会比较慢，
        #           而在有GPU的情况下，使用cuda可以显著提高模型的运行速度
        # 检查CUDA是否可用
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.logger.info(f"使用设备：{self.device}")
        # 重排序模型的路径
        rerank_model_path = os.path.join(rag_qa_path, "models", "bge-reranker-large")
        self.reranker = CrossEncoder(rerank_model_path, device=self.device, tokenizer_kwargs={
            "padding": "longest",  # 批量处理时自动pad到最长序列
            "truncation": True,  # 自动截断超长文本
            "max_length": 512,  # 重排序模型最优输入长度（BGE-Reranker默认）
            "return_tensors": "pt"
        })

        # 3.2：初始化BGE-M3嵌入模型（使用float16精度和cpu）
        # 参数说明：
        #   model_path: BGE-M3模型的路径，用于加载预训练的混合向量生成模型
        #   device：用来指定模型运行的设备，cpu运行。
        #   use_fp16：是否使用float16精度，这里设置为True，以减少内存占用，默认为False
        #           这里选择True是因为在没有GPU的情况下，也可以运行模型，只是速度会比较慢，
        #           而在有GPU的情况下，使用float16精度可以显著提高模型的运行速度
        # BGE-M3模型路径
        bge_m3_model_path = os.path.join(rag_qa_path, "models", "bge-m3")
        self.embedding_function = BGEM3EmbeddingFunction(bge_m3_model_path,
                                                         device=self.device, use_fp16=False)
        # 打印向量维度, 返回一个字典，包含三个键值对，分别是dense、colbert_vecs、sparse
        #   dense: 向量的维度，这里是1024，稠密向量维度
        #   colbert_vecs: 向量的维度，这里是1024，colbert向量维度，表示每个token被编码成的向量维度（1024），在混合检索中
        #                 colbert_vecs向量提供了介于稠密全局表示和稀疏词汇表示之间的细粒度匹配能力，能够提升检索的准确度和相关性
        #   sparse: 向量的维度，这里是250002，稀疏向量维度，稀疏向量词汇表的大小，词汇表大小表示模型能够处理的唯一词汇数量（250002个）
        # {'dense': 1024, 'colbert_vecs': 1024, 'sparse': 250002}
        print(f"embedding_function->device:{self.embedding_function.dim}")

        # 获取稠密向量的维度
        # 稠密向量是指在BGE-M3模型中生成的固定长度的向量，表示文本的语义信息
        # 这里通过调用模型dim属性获取稠密向量的维度，用于后续的向量存储和检索
        self.dense_dim = self.embedding_function.dim['dense']

        # 4：初始化Milvus客户端
        self.client = MilvusClient(uri=f"http://{self.host}:{self.port}", db_name=self.database)

        # 5：创建向量集合
        self._create_or_load_collection()

    """
    创建向量集合
    实现思路：
    1：判断向量集合是否存在
        1.1：如果不存在，则创建向量集合
            1.1.1：定义字段schema，包括id、向量字段、文本字段
            1.1.2：定义向量索引信息
            1.1.3：创建向量集合，指定字段schema
        1.2：如果存在，则直接使用该向量集合
    """

    def _create_or_load_collection(self):
        """
        创建向量集合
        :param client: Milvus客户端
        :param collection_name: 向量集合名称
        :param dimension: 向量维度
        :return: None
        """
        if not self.client.has_collection(self.collection_name):
            self.logger.info(f"向量集合 {self.collection_name} 不存在，正在创建...")
            # 创建集合的schema，禁用自动id，启动动态字段
            schema = self.client.create_schema(auto_id=False, enable_dynamic_field=True)
            # 添加id列。作为主键（varchar类型）
            schema.add_field(field_name="id", datatype=DataType.VARCHAR, max_length=100, is_primary=True)
            # 添加文本字段，VARCHAR 类型，最大长度 65535
            schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=65535)
            # 添加稠密向量字段，FLOAT_VECTOR 类型，维度由嵌入函数指定
            schema.add_field(field_name="dense_vector", datatype=DataType.FLOAT_VECTOR, dim=self.dense_dim)
            # 添加稀疏向量字段，SPARSE_FLOAT_VECTOR 类型
            schema.add_field(field_name="sparse_vector", datatype=DataType.SPARSE_FLOAT_VECTOR)
            # 添加父块 ID 字段，VARCHAR 类型，最大长度 100
            schema.add_field(field_name="parent_id", datatype=DataType.VARCHAR, max_length=100)
            # 添加父块内容字段，VARCHAR 类型，最大长度 65535
            schema.add_field(field_name="parent_content", datatype=DataType.VARCHAR, max_length=65535)
            # 添加学科类别字段，VARCHAR 类型，最大长度 50
            schema.add_field(field_name="source", datatype=DataType.VARCHAR, max_length=50)
            # 添加时间戳字段，VARCHAR 类型，最大长度 50
            schema.add_field(field_name="timestamp", datatype=DataType.VARCHAR, max_length=50)

            # 准备索引字段信息
            # 用于创建索引参数对象，后续添加索引时会使用到该对象
            index_params = self.client.prepare_index_params()

            # 为稠密向量字段添加IVF_FLAT索引
            index_params.add_index(
                field_name="dense_vector",  # 稠密向量字段名称
                index_name="dense_index",  # 索引名称
                index_type="IVF_FLAT",  # 倒排文件索引
                metric_type="IP",  # 内积相似度
                params={"nlist": 1024}  # 聚类中心数量
            )

            # 为稀疏向量字段添加SPARSE_INVERTED_INDEX索引
            index_params.add_index(
                field_name="sparse_vector",  # 稀疏向量字段名称
                index_name="sparse_index",  # 索引名称
                index_type="SPARSE_INVERTED_INDEX",  # 稀疏倒排索引
                metric_type="IP",  # 内积相似度
                params={"drop_ratio_build": 0.2}  # 构建时丢弃比例
            )
            # 创建向量集合
            self.client.create_collection(self.collection_name, schema=schema, index_params=index_params)
            self.logger.info(f"向量集合 {self.collection_name} 已创建...")
        else:
            self.logger.info(f"向量集合 {self.collection_name} 已存在，正在加载...")

    """
    向向量集合中添加文档
    :param documents: 文档列表
    :return: None
    """

    def add_documents(self, documents: list[Document]):
        # 将分块后的文档转换成向量并存储到milvus集合中
        # 获取文档内容
        document_contents = [doc.page_content for doc in documents]
        # 将获取到的文档内容进行嵌入转换向量化（稠密+稀疏）
        embeddings = self.embedding_function(document_contents)
        print(embeddings)

        # 准备存储向量数据的集合列表
        vectors = []

        # 遍历每个文档块
        for index, doc in enumerate(documents):
            # 找到文档块的内容，生成唯一的id（md5哈希）
            content = doc.page_content.encode("utf-8")
            id = hashlib.md5(content).hexdigest()
            print(f"稠密向量 id: {id}，稠密向量: {embeddings['dense'][index]}，数量为：{len(embeddings['dense'][index])}")
            # 处理稀疏向量，将其转换成字典格式，并添加到向量列表中
            sparse_vector = {}
            # 获取当前文档的稀疏向量矩阵
            # 参数：
            #   embeddings["sparse"]：所有文档的稀疏向量表示（coo格式）
            #   [[index], :]：选择当前文档的稀疏索引（行索引为index）
            # 返回值sparse_matrix：当前文档的稀疏向量矩阵（coo格式）
            """
                示例数据
                    # 原始稀疏矩阵（3个文档，词汇表大小为5）
                    # 文档1: [0, 0.5, 0, 0.3, 0] 
                    # 文档2: [0.2, 0, 0, 0, 0.7]
                    # 文档3: [0, 0, 0.4, 0, 0]

                    # COO格式表示:
                    embeddings["sparse"].indices = [1, 3, 0, 4, 2]  # 非零元素列索引
                    embeddings["sparse"].data = [0.5, 0.3, 0.2, 0.7, 0.4]  # 非零元素值
                    embeddings["sparse"].shape = (3, 5)  # 3行5列

                    # 当index=1时:
                    embeddings["sparse"][[1], :]  # 提取第2个文档的稀疏向量
                    # 结果: 一个只包含[0.2, 0, 0, 0, 0.7]的稀疏矩阵
                    # indices=[0, 4], data=[0.2, 0.7]
            """
            sparse_matrix = embeddings["sparse"][[index], :]
            print(sparse_matrix)
            # 获取稀疏矩阵中非零元素的列索引位置
            indices = sparse_matrix.indices
            # 获取稀疏矩阵中非零元素对应的数值
            value = sparse_matrix.data

            # 根据非零元素的索引和值构建稀疏向量矩阵
            for token_id, value in zip(indices, value):
                sparse_vector[token_id] = value

            print(f"稀疏向量 id: {id}，稀疏向量: {sparse_vector}，数量为：{len(sparse_vector)}")
            vectors.append({
                "id": id,  # 文档id
                "text": doc.page_content,  # 文档内容
                "dense_vector": embeddings["dense"][index],  # 稠密向量
                "sparse_vector": sparse_vector,  # 稀疏向量
                "parent_id": doc.metadata["parent_id"],  # 父块id
                "parent_content": doc.metadata["parent_content"],  # 父块内容
                "source": doc.metadata["source"],  # 文档来源
                "timestamp": doc.metadata["timestamp"],  # 时间戳
            })

        # 向向量集合中插入数据
        if vectors:
            self.client.upsert(collection_name=self.collection_name, data=vectors)
            self.logger.info(f"已向向量集合 {self.collection_name} 中添加 {len(vectors)} 条文档向量...")
        else:
            self.logger.warning("没有需要添加的文档向量...")

    """
    实现混合检索并重新排序，返回最相关的K个文档块
    :param query: 查询文本
    :param source_filter: 来源筛选（可选）
    :param k: 返回的文档块数量（默认5）
    :return: 最相关的K个文档块列表
    """

    def hybrid_search_with_rerank(self, query: str, k=config.RETRIEVAL_K, source_filter=None):
        # 将查询的文本字符串转换为向量表示
        query_embedding = self.embedding_function([query])

        # 获取到查询文本的稠密向量表示
        query_dense_vector = query_embedding["dense"][0]

        # 定义处理查询的稀疏向量矩阵
        sparse_query_vector = {}
        # 获取查询文本的稀疏向量表示
        query_sparse_vector = query_embedding["sparse"][[0], :]
        # 获取稀疏向量的非零元素索引
        query_sparse_indices = query_sparse_vector.indices
        # 获取稀疏向量的非零元素值
        query_sparse_values = query_sparse_vector.data

        # 构建稀疏向量矩阵表示
        for token_id, value in zip(query_sparse_indices, query_sparse_values):
            sparse_query_vector[token_id] = value

        print(f"查询文本的稠密向量表示: {query_dense_vector}，数量为：{len(query_dense_vector)}")
        print(f"查询文本的稀疏向量表示: {query_sparse_vector}，数量为：{len(query_sparse_indices)}")

        # 构建数据来源过滤表达式（如果有）
        filter_expr = f'source=="{source_filter}"' if source_filter else ''

        # 构建稠密向量查询请求
        dense_request = AnnSearchRequest(
            data=[query_dense_vector],  # 查询的稠密向量表示
            anns_field="dense_vector",  # 稠密向量字段名称
            param={'metric_type': 'IP', 'params': {'nprobe': 10}},  # 相似度计算指标（内积）
            limit=k,  # 返回的文档块数量
            expr=filter_expr,  # 来源过滤表达式
        )

        # 构建稀疏向量查询请求
        sparse_request = AnnSearchRequest(
            data=[query_sparse_vector],  # 稀疏向量表示
            anns_field="sparse_vector",  # 稀疏向量字段名称
            param={'metric_type': 'IP', 'params': {}},  # 相似度计算指标（内积相似度）
            limit=k,  # 返回的文档块数量
            expr=filter_expr,  # 来源过滤表达式
        )

        # 创建混合权重排序器（稠密1.0, 稀疏0.7）
        ranker = WeightedRanker(1.0, 0.7)

        # 执行混合搜索，结合稠密向量和稀疏向量进行检索
        # 参数说明:
        #   collection_name: 指定要搜索的集合名称
        #   reqs: 包含稠密向量查询请求和稀疏向量查询请求的列表
        #   ranker: 用于对搜索结果进行排序的排序器
        #   limit: 返回结果的数量限制
        #   output_fields: 指定需要返回的字段列表
        # 返回值: 搜索结果列表，这里取第一个结果
        results = self.client.hybrid_search(
            collection_name=self.collection_name,
            reqs=[dense_request, sparse_request],  # 稠密向量查询请求和稀疏向量查询请求
            ranker=ranker,
            limit=k,
            output_fields=["text", "parent_id", "parent_content", "source", "timestamp"]
        )[0]

        # print(f"混合检索结果: {results}")

        # 将搜索结果转换为document对象
        sub_chunks = [self._doc_from_hit(hit['entity']) for hit in results]
        # print(f"子块列表: {sub_chunks}")
        # 获取唯一的父文档（去重）
        parent_docs = self._get_unique_parent_docs(sub_chunks)

        # print(f"父块列表: {parent_docs}")

        # 当有多个父文档时进行重排序
        if parent_docs and len(parent_docs) > 1:
            # 创建查询-文档对
            query_doc_pairs = [(query, doc.page_content) for doc in parent_docs]

            # 使用重排序模型预测相关性分数
            scores = self.reranker.predict(query_doc_pairs)
            # print(f"相关性分数: {scores}")

            # 创建（分数，文档）对
            score_doc_pairs = list(zip(scores, parent_docs))

            # 按照分数降序排序
            sorted_pairs = sorted(score_doc_pairs, key=lambda x: x[0], reverse=True)
            # print(f"排序后的（分数，文档）对: {sorted_pairs}")

            # 提取排序后的文档
            sorted_docs = [doc for score, doc in sorted_pairs]
            randed_parent_docs = sorted_docs
        else:
            randed_parent_docs = parent_docs

        # 返回前M个候选文档
        return randed_parent_docs[:config.CANDIDATE_M]

    # 定义私有方法，从子块中提取去重的父文档
    def _get_unique_parent_docs(self, sub_chunks):
        # 初始化集合，用于存储已处理的父块内容（去重）
        parent_contents = set()
        # 初始化列表，用于存储唯一父文档
        unique_docs = []
        # 遍历所有子块
        for chunk in sub_chunks:
            # 获取子块的父块内容，默认为子块内容
            parent_content = chunk.metadata.get("parent_content", chunk.page_content)
            # 检查父块内容是否非空且未重复
            if parent_content and parent_content not in parent_contents:
                # 创建新的 Document 对象，包含父块内容和元数据
                unique_docs.append(Document(page_content=parent_content, metadata=chunk.metadata))
                # 将父块内容添加到去重集合
                parent_contents.add(parent_content)
        # 返回去重后的父文档列表
        return unique_docs

    # 将hit转换成document对象
    def _doc_from_hit(self, hit):
        return Document(
            page_content=hit["text"],  # 文档内容
            metadata={
                'source': hit['source'],  # 文档来源
                'timestamp': hit['timestamp'],  # 时间戳
                'parent_id': hit['parent_id'],  # 父块id
                'parent_content': hit['parent_content'],  # 父块内容
            }
        )


if __name__ == '__main__':
    # 向量存储实例化
    vector_store = VectorStore()

    # # 处理文档并添加到向量集合中
    documents = process_documents("F://code//integrated_qa_system//rag_qa//data//ai_data")

    # 添加文档到向量集合中
    vector_store.add_documents(documents)
    query = "AI学科的课程内容是什么?"
    results = vector_store.hybrid_search_with_rerank(query, source_filter="ai")
    print(f"重排序后的文档: {results}")
    print(f"文档数量: {len(results)}")

"""
稠密向量(Dense Vector)和稀疏向量(Sparse Vector)混合搜索的主要优势体现在以下几个方面：

1. 互补的语义表示能力
   - 稠密向量能够捕捉词与词之间的语义相似性，对于同义词、近义词具有良好的泛化能力
   - 稠密向量通过神经网络训练得到，在高维空间中语义相近的词汇距离较近
   - 稀疏向量保留了精确的词汇匹配信息，对关键词出现与否非常敏感
   - 稀疏向量通常基于词频或者BM25等传统检索方法，能准确反映词汇的重要性

2. 提升检索准确性
   - 使用了加权排序器 `WeightedRanker(1.0，0.7 )` 对两种向量的结果进行融合
   - 稠密向量权重设为1.0，稀疏向量权重设为0.7，这样既考虑了语义相似性又兼顾了关键词匹配
   - 可以同时捕获查询和文档间的语义关联以及精确的词汇匹配关系

3. 增强鲁棒性和适应性
   - 当查询中的关键词在文档中没有完全匹配时，稠密向量可以通过语义相似性找到相关内容
   - 当存在专业术语或特定词汇时，稀疏向量可以确保这些重要词汇被准确匹配
   - 这种组合方式比单独使用任一种向量都更加稳健，适用于各种不同的查询场景

4. 优化搜索性能
   - 在 Milvus 中分别建立了针对 `dense_vector` 和 `sparse_vector` 的索引：
     - 稠密向量采用 IVF_FLAT 索引，适合高维稠密向量的快速近似最近邻搜索
     - 稀疏向量采用 SPARSE_INVERTED_INDEX 索引，专门针对稀疏向量设计，提高检索效率
   - 通过 `AnnSearchRequest` 并行执行两次搜索，并最终合并结果，提高了整体搜索速度

5. 更好的上下文理解
   - 提取唯一父文档并在必要时利用 `CrossEncoder` 进行重排序
   - 结合父文档内容进行综合判断，避免因局部片段导致的误判
   - 这有助于更全面地理解文档内容及其与查询的相关性

综上所述，这种混合搜索策略充分发挥了稠密向量和稀疏向量各自的优势，提供了更为精准且高效的检索体验。
"""


